# 使用agent代理连接数据库


import os
from operator import itemgetter

from langchain.agents import AgentType, create_openai_functions_agent, AgentExecutor
from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.agent_toolkits import SQLDatabaseToolkit, create_sql_agent
from langchain_community.tools import QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_deepseek import ChatDeepSeek
from langgraph.prebuilt import chat_agent_executor

# 本地有clash代理，所以配置一下，不然一些依赖下载不下来（chroma）,或者将代理关闭后下载
os.environ['http_proxy'] = '127.0.0.1:7890'
os.environ['https_proxy'] = '127.0.0.1:7890'

os.environ["LANGCHAIN_TRACING_V3"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"
os.environ["LANGCHAIN_PROJECT"] = "langchain-community-demo"

deepseek_api_key = "sk-1dd16a258a73428d910d38c782e1c94f"

# 1、创建模型
model = ChatDeepSeek(
    api_key=deepseek_api_key,
    # deepseek-reasoner : DeepSeek-R1
    # deepseek-chat : DeepSeek-V3
    model="deepseek-chat",
    temperature=0.7,
    max_tokens=1024
)

# 本例以 SQLAlchemy 演示
HOSTNAME = "127.0.0.1"
PORT = "3306"
DATABASE = "book_system"
USERNAME = "root"
PASSWORD = "123456"
# mysqlclient驱动URL
# MYSQL_URL = f"mysql+mysqldb://{USERNAME}:{PASSWORD}@{HOSTNAME}:{PORT}/{DATABASE}?charset=utf8mb4"
MYSQL_URL = 'mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)

db = SQLDatabase.from_uri(MYSQL_URL)

# 创建工具
toolkit = SQLDatabaseToolkit(db=db, llm=model)
tools = toolkit.get_tools()

# 使用agent完成整个数据库的整合
system_prompt = """
你是一个被设计用来与SQL数据库交互的代理。
给定一个输入问题，创建一个语法正确的SQL语句并执行，然后查看查询结果并返回答案。
除非用户指定了他们想要获得的示例的具体数量，否则始终将SQL查询数量限制为最多10个结果。
你可以按相关列对结果进行排序，以返回MySQL数据库中做匹配的数据。
你可以使用与数据库交互的工具，在执行查询之前，你必须仔细检查，如果在执行时出现错误，请重写查询SQL并重试。
不要对数据库做任何DML语句（插入，更新，删除等）。

首先 ，你应该查看数据库中的表，看看可以查询什么。
不要跳过这一步。
然后查询最相关表的模式。
"""

# system_message = PromptTemplate.from_template(system_prompt)
system_message = SystemMessage(content=system_prompt)

# 创建代理（不同模型支持的调用方式不同 第一种openai ,第二种deepseek） deepseek不支持函数调用
agent_executor = chat_agent_executor.create_tool_calling_executor(model, tools, prompt=system_message)
# agent_executor = create_sql_agent(llm=model, toolkit=toolkit, verbose=True,
#                                   agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION)

# prompt = ChatPromptTemplate(system_prompt)
# agent = create_openai_functions_agent(model, tools, prompt)
# agent_executor = AgentExecutor(agent, tools, verbose=True)

# # 可循环提问
# questions = [
#     "书籍类型表中多少条数据？",
#     "书籍类型表中有哪些书籍类型？"
# ]
# for question in questions:
#     resp = agent_executor.invoke({'input': question})
#     print(resp['output'])


resp = agent_executor.invoke({'messages': [HumanMessage(content='书籍类型表中多少条数据？')]})
print('最后的拉拉==>', resp)
result = resp['messages']
print('最后的resp===>', result)
print(len(result))
print(result[len(result) - 1])